module lib_module

    implicit none

    private
    public :: nx, ny, compute_Mandebrot, compute_Mandebrot_

    type complex_type
        real :: re
        real :: im
    end type complex_type

    integer, parameter :: max_iter = 200
    integer, parameter :: nx = 960, ny = 768
    real, parameter :: xmin = -2.0, xmax = 0.6, ymin = -1.5, ymax = 1.5

contains

    !> 求复数的模的平方
    pure function norm2(z) result(r)
        type(complex_type), intent(in) :: z
        real :: r

        r = z%re**2 + z%im**2

    end function norm2

    !> 复数加法
    pure function add(z1, z2) result(z)
        type(complex_type), intent(in) :: z1, z2
        type(complex_type) :: z

        z%re = z1%re + z2%re
        z%im = z1%im + z2%im

    end function add

    !> 复数乘法
    pure function mul(z1, z2) result(z)
        type(complex_type), intent(in) :: z1, z2
        type(complex_type) :: z

        z%re = z1%re*z2%re - z1%im*z2%im
        z%im = z1%re*z2%im + z1%im*z2%re

    end function mul

    !> mandelbrot kernel
    pure integer function mandelbrot_kernel(z) result(i)
        type(complex_type), intent(in) :: z
        type(complex_type) :: z0

        z0 = z
        do i = 1, max_iter
            z0 = add(mul(z0, z0), z)
            if (norm2(z0) > 4.0) return
        end do

    end function mandelbrot_kernel

    !> mandelbrot kernel (#)
    pure integer function mandelbrot_kernel_(z) result(i)
        complex, intent(in) :: z
        complex :: z0

        z0 = z
        do i = 1, max_iter
            z0 = z0*z0 + z
            if ((z0%re**2 + z0%im**2) > 4.0) return
        end do

    end function mandelbrot_kernel_

    !> mandelbrot set computation
    pure function compute_Mandebrot() result(ans)
        integer :: ans(nx, ny)
        real :: dx, dy, x, y
        integer :: i, j

        dx = (xmax - xmin)/nx
        dy = (ymax - ymin)/ny

        y = ymin
        do j = 1, ny
            x = xmin
            do i = 1, nx
                ans(i, j) = mandelbrot_kernel(complex_type(x, y))
                x = x + dx
            end do
            y = y + dy
        end do

    end function compute_Mandebrot

    !> mandelbrot set computation (#)
    pure function compute_Mandebrot_() result(ans)
        integer :: ans(nx, ny)
        real :: dx, dy, x, y
        integer :: i, j

        dx = (xmax - xmin)/nx
        dy = (ymax - ymin)/ny

        y = ymin
        do j = 1, ny
            x = xmin
            do i = 1, nx
                ans(i, j) = mandelbrot_kernel_(cmplx(x, y))
                x = x + dx
            end do
            y = y + dy
        end do

    end function compute_Mandebrot_

end module lib_module

program main

    use lib_module, only: nx, ny, compute_Mandebrot, compute_Mandebrot_
    implicit none

    integer :: ans(nx, ny), i
    real :: t1, t2

    call cpu_time(t1)
    do i = 1, 10
        ans = compute_Mandebrot()
    end do
    call cpu_time(t2)

    print "(a,t40,f0.1)", "drived type, time(ms): ", (t2 - t1)*1000.0/10

    call cpu_time(t1)
    do i = 1, 10
        ans = compute_Mandebrot_()
    end do
    call cpu_time(t2)

    print "(a,t40,f0.1)", "native complex, time(ms): ", (t2 - t1)*1000.0/10

end program main
